import gensim
from gensim import corpora, models
import numpy as np

class ladAnalysis:

    '''

    :parameter

    input : 输入数据
    modelTrainingFlag ： 是否开启LDA模型训练

    '''

    def __init__(self, input, modelTrainingFlag = False):
        self.modelTrainingFlag = modelTrainingFlag
        self.input = input
        self.dictionaryWords = np.load(r"D:\pythonProject\endOfTerm\word_list.npy", allow_pickle=True)

    def run(self):
        if self.modelTrainingFlag:
            self.modelTraining()
        else:
            self.modelAnalyze()

    #模型训练
    def modelTraining(self):
        dictionary = corpora.Dictionary(self.input)
        dictionary.save('dictionary.dictionary')
        corpus = [dictionary.doc2bow(doc) for doc in self.input]
        corpora.MmCorpus.serialize('corpus.mm', corpus)#保存词袋向量
        lda = gensim.models.ldamodel.LdaModel(corpus, num_topics=20, id2word=dictionary,passes=100, random_state=1)  # 调用LDA模型，20个主题；训练语料库50次
        lda.save(r"C:\Users\chenxu\Desktop\组会\LDA主题模型\LDA.model")

    #模型分析
    def modelAnalyze(self):
        lda = gensim.models.ldamodel.LdaModel.load(r"LDA.model")#加载已经训练号的lda模型
        dictionary = corpora.Dictionary(self.input)
        corpus = [dictionary.doc2bow(doc) for doc in self.input]
        lda.get_document_topics(corpus, minimum_probability=1e-8)
        # 输出对应文档的最大可能性的主题归类
        maxProbability = []
        maxTopic = []
        for zip in [*lda.get_document_topics(corpus, minimum_probability=1e-8)]:
            index = []
            probability = []
            for k in zip:
                index.append(k[0])
                probability.append(k[1])
            maxProbability.append(max(probability))
            maxTopic.append(probability.index(max(probability)))
        print(maxProbability)
        print(maxTopic)
        for i in maxTopic:
            print("LDA模型主题分布：", lda.print_topics(num_topics=30, num_words=5)[i])
